#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <stdbool.h>
#include <libgen.h>
#include "bl_crc32.h"
#include "package.h"

#include "tinycrypt_config.h"
#include "tinycrypt.h"
#include "fastlz.h"
#include "quicklz.h"

/* compress level: 1 or 2 */
#define FASTLZ_COMPRESS_LEVEL          2

#define BLOCK_HEADER_SIZE              4

#define BOOT_ENCRYP_AES_IV  	"CRITICALERROR!!!"
#define BOOT_ENCRYP_AES_KEY 	"CRITICALERROR!!!REBOOT SYSTEM!!!"

char *outfile = "update.tbl"; /* 输出文件的文件名 */

typedef struct {
    uint32_t magic_numb;
    uint32_t tbl_ver;        // [31:24] major, [23:12] minor [11:0] reversion
    char model[8];
    uint16_t flag;
    uint16_t block_size;
    uint32_t build_ts;
    uint32_t ctx_crc;
    uint32_t raw_size;
    uint32_t ctx_size;
    char parti_name[12];
    char reserve1[16];
    char reserve2[12];
    uint32_t header_crc;
}bl_header_info, *bl_header_info_t;

static void printUsage(const char *appname)
{
    printf("Usage: %s \r\n"
            "        [--aes256 || -a]                           Use aes256 encryption[N]\r\n"
            "        [[--quicklz || -q] || [--fastlz || -f]]    Use compression quicklz or fastlz[N]\r\n"
            "        [--bs || -b] <compression block size>      512 1024 2048 4096max[Default]\r\n"
            "        --model || -m <instrument model string>    8max\r\n"
            "        --partition || -p <flash partition>        12max\r\n"
            "        -i <src bin>                               Input bin file\r\n"
            "        [-o <out path>]                            Output file, update.tbl default\r\n"
            , basename(appname));
}

static int test_tiny_aes(void)
{
    tiny_aes_context ctx;
    uint8_t iv[16 + 1];
    uint8_t private_key[32 + 1];

    unsigned char data[] = "1234567890123456";
    unsigned char data_encrypt[32];
    unsigned char data_decrypt[32];

    /* encrypt */
    memcpy(iv, BOOT_ENCRYP_AES_IV, strlen(BOOT_ENCRYP_AES_IV));
    iv[sizeof(iv) - 1] = '\0';
    memcpy(private_key, BOOT_ENCRYP_AES_KEY, strlen(BOOT_ENCRYP_AES_KEY));
    private_key[sizeof(private_key) - 1] = '\0';

    memset(data_encrypt, 0x0, sizeof(data_encrypt));
    tiny_aes_setkey_enc(&ctx, (uint8_t *) private_key, 256);
    tiny_aes_crypt_cbc(&ctx, AES_ENCRYPT, strlen(data), iv, data, data_encrypt);

    /* decrypt */
    memcpy(iv, BOOT_ENCRYP_AES_IV, strlen(BOOT_ENCRYP_AES_IV));
    iv[sizeof(iv) - 1] = '\0';
    memcpy(private_key, BOOT_ENCRYP_AES_KEY, strlen(BOOT_ENCRYP_AES_KEY));
    private_key[sizeof(private_key) - 1] = '\0';

    memset(data_decrypt, 0x0, sizeof(data_decrypt));
    tiny_aes_setkey_dec(&ctx, (uint8_t *) private_key, 256);
    tiny_aes_crypt_cbc(&ctx, AES_DECRYPT, strlen(data), iv, data_encrypt, data_decrypt);

    if(memcmp(data, data_decrypt, strlen(data)) == 0)
    {
        printf("AES CBC mode passed!\n");
        return 0;
    }
    else
    {
        printf("AES CBC mode failed!");
        return -1;
    }
}

int tboot_archive(int argc, char *argv[])
{
    int arg_c = argc-1;
    char **arg_v = &argv[1];
    bool aes256_en = false, quicklz_en = false, fastlz_en = false;
    char *model = NULL, *partition = NULL, *ifile = NULL, *ofile = NULL;
    size_t block_sz = BLOCK_SIZE;

//    test_tiny_aes();

    if(argc < 2){
        printUsage(argv[0]);
        return 0;
    }

    while(arg_c > 0){
        if(strcmp(arg_v[0], "--aes256") == 0 || strcmp(arg_v[0], "-a") == 0){
            aes256_en = true;
        }else if(strcmp(arg_v[0], "--quicklz") == 0 || strcmp(arg_v[0], "-q") == 0){
            quicklz_en = true;
        }else if(strcmp(arg_v[0], "--fastlz") == 0 || strcmp(arg_v[0], "-f") == 0){
            fastlz_en = true;
        }else if(strcmp(arg_v[0], "--model") == 0 || strcmp(arg_v[0], "-m") == 0){
            model = arg_v[1];
        }else if(strcmp(arg_v[0], "--partition") == 0 || strcmp(arg_v[0], "-p") == 0){
            partition = arg_v[1];
        }else if(strcmp(arg_v[0], "--bs") == 0 || strcmp(arg_v[0], "-b") == 0){
            block_sz = atoi(arg_v[1]);
        }else if(strcmp(arg_v[0], "-i") == 0){
            ifile = arg_v[1];
        }else if(strcmp(arg_v[0], "-o") == 0){
            ofile = arg_v[1];
        }else{
        }
        arg_c--; arg_v++;
    }
    if (((model == NULL) || (strlen(model) > 8)) ||
        ((partition == NULL) || (strlen(partition) > 12)) ||
        (quicklz_en && fastlz_en) ||
        (ifile == NULL) ||
        (block_sz != 512 && block_sz != 1024 && block_sz != 2048 && block_sz != 4096)) {
        printUsage(argv[0]);
        return 0;
    }

    if (ofile != NULL) {
        outfile = ofile;
    }

    FILE *ifp, *ofp;
    struct stat st;
    time_t tim;
    bl_header_info header;
    size_t raw_size = 0;
    size_t ctx_size = 0;
    uint8_t *buf = NULL, *compr_buf = NULL, *decrypt_buf = NULL, *encrypt_buf = NULL, *ctx_buf = NULL;
    size_t padding_sz = 0, compr_sz = 0, decrypt_sz = 0, encrypt_sz = 0;
    /* Start to compress file  */
    qlz_state_compress *state_compress = NULL;
    uint8_t compr_hdr[BLOCK_HEADER_SIZE] = {0};
    tiny_aes_context *aes_ctx = NULL;
    uint8_t *aes_iv = NULL;
    size_t rd;

    if(stat(ifile, &st) < 0){
        perror("Stat file ");
        return -1;
    }
    raw_size = st.st_size;
printf("raw size: %d\n", raw_size);
    ifp = fopen(ifile, "rb");
    if(!ifp){
        printf("Open file %s failure!\n", ifile);
        return -1;
    }
    ofp = fopen(outfile, "wb+");
    if(!ofp){
        printf("Open file %s failure!\n", outfile);
        fclose(ifp);
        return -1;
    }

    header.magic_numb = IDENTIFIER;
    header.tbl_ver = SOFTWARE_VERSION;
    strncpy(header.model, model, 8);
    header.flag = aes256_en | ((quicklz_en ? 2 : fastlz_en ? 3 : 0) << 8);
    header.block_size = BLOCK_SIZE;
    strncpy(header.parti_name, partition, 12);

    //time(&tim);
    header.build_ts = st.st_ctime;//tim;
    header.raw_size = raw_size;

    buf = malloc(header.block_size);
    ctx_buf = malloc(header.block_size * 1024);
    decrypt_buf = malloc(header.block_size * 1024);
    // aes256
    if (aes256_en) {
        aes_ctx = malloc(sizeof(tiny_aes_context));
        if (aes_ctx == NULL) {
            goto _err_exit;
        }
        aes_iv = malloc(17);
        if (aes_iv == NULL) {
            goto _err_exit;
        }

        memcpy(aes_iv, BOOT_ENCRYP_AES_KEY, 16);
        aes_iv[16] = 0;
        tiny_aes_setkey_enc(aes_ctx, BOOT_ENCRYP_AES_KEY, 256);
        encrypt_buf = malloc(header.block_size + padding_sz);
    }
    // compress
    if (fastlz_en) {
        padding_sz = FASTLZ_BUFFER_PADDING(header.block_size);
        compr_buf = malloc(header.block_size + padding_sz);
    } else if (quicklz_en) {
        padding_sz = QLZ_BUFFER_PADDING;
        state_compress = (qlz_state_compress *) malloc(sizeof(qlz_state_compress));
        if (!state_compress)
        {
            return -1;
        }
        memset(state_compress, 0x00, sizeof(qlz_state_compress));
        compr_buf = malloc(header.block_size + padding_sz);
    }
    while (1) {
        rd = fread(buf, 1, header.block_size, ifp);
        if (rd <= 0) break;
        if (fastlz_en || quicklz_en) {
            if (fastlz_en) {
                compr_sz = fastlz_compress_level(FASTLZ_COMPRESS_LEVEL, buf, rd, (char *)compr_buf);
            } else if (quicklz_en) {
                compr_sz = qlz_compress(buf, (char *)compr_buf, rd, state_compress);
            }
            compr_hdr[0] = (compr_sz >> 24) & 0xFF;
            compr_hdr[1] = (compr_sz >> 16) & 0xFF;
            compr_hdr[2] = (compr_sz >> 8) & 0xFF;
            compr_hdr[3] = compr_sz & 0xFF;
            memcpy(decrypt_buf + decrypt_sz, compr_hdr, BLOCK_HEADER_SIZE);
            decrypt_sz += BLOCK_HEADER_SIZE;
            memcpy(decrypt_buf + decrypt_sz, compr_buf, compr_sz);
            decrypt_sz += compr_sz;
        } else {
            memcpy(decrypt_buf + decrypt_sz, buf, rd);
            decrypt_sz += rd;
        }
    }

    // encryption
    if (aes256_en) {
        while (encrypt_sz < decrypt_sz) {
            if ((encrypt_sz + BLOCK_SIZE) <= decrypt_sz) {
                block_sz = BLOCK_SIZE;
            } else {
                block_sz = decrypt_sz - encrypt_sz;
            }
            memcpy(buf, decrypt_buf + encrypt_sz, block_sz);
            tiny_aes_crypt_cbc(aes_ctx, AES_ENCRYPT, block_sz, aes_iv, buf, encrypt_buf);
            encrypt_sz += block_sz;
            block_sz = (block_sz + 15) & (~0x0F);
            memcpy(ctx_buf + ctx_size, encrypt_buf, block_sz);
            ctx_size += block_sz;
        }
    } else {
        memcpy(ctx_buf, decrypt_buf, decrypt_sz);
        ctx_size += decrypt_sz;
    }

    // context crc & size
    header.ctx_crc = bl_crc32sum(ctx_buf, ctx_size);
    header.ctx_size = ctx_size;
printf("ctx size: %d\n", ctx_size);
printf("ctx crc: %x\n", header.ctx_crc);
    // header crc
    header.header_crc = bl_crc32sum(&header, sizeof(bl_header_info)-4);
printf("hdr crc: %x\n", header.header_crc);
    // write out file
    fwrite(&header, sizeof(bl_header_info), 1, ofp);
    fwrite(ctx_buf, ctx_size, 1, ofp);

_err_exit:
    if (buf) free(buf);
    if (ctx_buf) free(ctx_buf);
    if (decrypt_buf) free(decrypt_buf);
    if (aes_ctx) free(aes_ctx);
    if (aes_iv) free(aes_iv);
    if (state_compress) free(state_compress);
    if (compr_buf) free(compr_buf);
    if (encrypt_buf) free(encrypt_buf);

    fclose(ofp);
    fclose(ifp);

    return 0;
}
